import torch.optim
from torch.nn import Module


class Server:
    def __init__(self,
                 model: Module,
                 lr: float = 0.1,
                 optimizer: str = "SGD",
                 weight_decay: float = 1e-4,
                 momentum: float = 0.9,
                 scheduler: str = "steplr",
                 step_size: int = 30,
                 gamma: float = 0.1,
                 ) -> None:
        self.model = model
        self.optimizer = None
        if optimizer == "SGD":
            self.optimizer = torch.optim.SGD(self.model.parameters(), lr=lr, weight_decay=weight_decay,
                                             momentum=momentum)
        else:
            pass
        self.scheduler = None
        if scheduler == "steplr":
            self.scheduler = torch.optim.lr_scheduler.StepLR(self.optimizer, step_size, gamma)
        else:
            pass
